"""
向量存储服务（使用FAISS）
"""
import os
import pickle
import numpy as np
from typing import List, Tuple, Optional
import logging

logger = logging.getLogger(__name__)

try:
    import faiss
    FAISS_AVAILABLE = True
except ImportError:
    logger.warning("FAISS未安装，向量检索功能不可用")
    FAISS_AVAILABLE = False


def get_embedding_dimension(model_name: str) -> int:
    """
    根据模型名称返回向量维度
    
    Args:
        model_name: Ollama 模型名称
        
    Returns:
        向量维度
    """
    # 移除 :latest 等标签
    model_base = model_name.split(':')[0].lower()
    
    # 已知模型的维度映射
    dimension_map = {
        'nomic-embed-text': 768,
        'bge-m3': 1024,
        'bge-large-zh-v1.5': 1024,
        'bge-base-zh-v1.5': 768,
        'mxbai-embed-large': 1024,
    }
    
    return dimension_map.get(model_base, 1024)  # 默认 1024


class VectorStore:
    """向量存储和检索"""
    
    def __init__(self, dimension: Optional[int] = None):
        """
        初始化向量存储
        
        Args:
            dimension: 向量维度，如果为 None 则从配置中自动推断
        """
        if dimension is None:
            # 从配置中读取模型并自动推断维度
            try:
                from app.core.config import settings
                dimension = get_embedding_dimension(settings.OLLAMA_EMBED_MODEL)
                logger.info(f"自动推断向量维度: {dimension} (模型: {settings.OLLAMA_EMBED_MODEL})")
            except Exception as e:
                logger.warning(f"无法自动推断维度，使用默认值 1024: {e}")
                dimension = 1024
        self.dimension = dimension
        self.index = None
        self.texts = []
        self.metadatas = []
        
        if FAISS_AVAILABLE:
            # 使用FAISS的IndexFlatIP（内积索引）配合归一化向量实现余弦相似度
            # 余弦相似度 = 归一化向量的内积
            self.index = faiss.IndexFlatIP(dimension)
    
    def add_texts(
        self,
        texts: List[str],
        embeddings: List[List[float]],
        metadatas: Optional[List[dict]] = None
    ):
        """
        添加文本和向量
        
        Args:
            texts: 文本列表
            embeddings: 向量列表
            metadatas: 元数据列表
        """
        if not FAISS_AVAILABLE:
            logger.error("FAISS未安装，无法添加向量")
            return
        
        try:
            # 转换为numpy数组
            vectors = np.array(embeddings, dtype=np.float32)
            
            # 归一化向量（用于余弦相似度）
            faiss.normalize_L2(vectors)
            
            # 添加到FAISS索引
            self.index.add(vectors)
            
            # 保存文本和元数据
            self.texts.extend(texts)
            if metadatas:
                self.metadatas.extend(metadatas)
            else:
                self.metadatas.extend([{}] * len(texts))
            
            logger.info(f"成功添加 {len(texts)} 个文本到向量库")
        except Exception as e:
            logger.error(f"添加向量失败: {str(e)}")
            raise
    
    def similarity_search(
        self,
        query_embedding: List[float],
        k: int = 5
    ) -> List[Tuple[str, float, dict]]:
        """
        相似度搜索
        
        Args:
            query_embedding: 查询向量
            k: 返回top-k结果
            
        Returns:
            (文本, 相似度分数, 元数据) 的列表
        """
        if not FAISS_AVAILABLE or self.index is None:
            logger.error("FAISS未安装或索引未初始化")
            return []
        
        try:
            # 转换为numpy数组
            query_vector = np.array([query_embedding], dtype=np.float32)
            
            # 归一化查询向量（用于余弦相似度）
            faiss.normalize_L2(query_vector)
            
            # 搜索（IndexFlatIP返回的是内积，归一化后等于余弦相似度）
            scores, indices = self.index.search(query_vector, k)
            
            results = []
            for i, (score, idx) in enumerate(zip(scores[0], indices[0])):
                if idx < len(self.texts):
                    # IndexFlatIP返回的是余弦相似度（范围-1到1）
                    # 转换为0-1范围：(score + 1) / 2
                    # 但通常score已经在0-1范围内，直接使用
                    similarity = float(max(0.0, min(1.0, score)))  # 确保在0-1范围
                    results.append((
                        self.texts[idx],
                        similarity,
                        self.metadatas[idx]
                    ))
            
            return results
        except Exception as e:
            logger.error(f"向量搜索失败: {str(e)}")
            return []
    
    def save(self, filepath: str):
        """
        保存向量库到文件
        
        Args:
            filepath: 保存路径
        """
        if not FAISS_AVAILABLE:
            logger.error("FAISS未安装，无法保存")
            return
        
        try:
            os.makedirs(os.path.dirname(filepath), exist_ok=True)
            
            # 保存FAISS索引
            faiss.write_index(self.index, f"{filepath}.index")
            
            # 保存文本和元数据
            with open(f"{filepath}.meta", "wb") as f:
                pickle.dump({
                    "texts": self.texts,
                    "metadatas": self.metadatas,
                    "dimension": self.dimension
                }, f)
            
            logger.info(f"向量库已保存到 {filepath}")
        except Exception as e:
            logger.error(f"保存向量库失败: {str(e)}")
            raise
    
    def load(self, filepath: str):
        """
        从文件加载向量库
        
        Args:
            filepath: 文件路径
        """
        if not FAISS_AVAILABLE:
            logger.error("FAISS未安装，无法加载")
            return
        
        try:
            # 加载FAISS索引
            self.index = faiss.read_index(f"{filepath}.index")
            
            # 加载文本和元数据
            with open(f"{filepath}.meta", "rb") as f:
                data = pickle.load(f)
                self.texts = data["texts"]
                self.metadatas = data["metadatas"]
                self.dimension = data["dimension"]
            
            logger.info(f"向量库已从 {filepath} 加载，共 {len(self.texts)} 条")
        except Exception as e:
            logger.error(f"加载向量库失败: {str(e)}")
            raise
    
    def clear(self):
        """清空向量库"""
        if FAISS_AVAILABLE:
            self.index = faiss.IndexFlatL2(self.dimension)
        self.texts = []
        self.metadatas = []
        logger.info("向量库已清空")
    
    def count(self) -> int:
        """返回向量数量"""
        return len(self.texts)
    
    def remove_by_document_id(self, document_id: str):
        """
        删除指定文档的所有向量
        
        Args:
            document_id: 文档ID
            
        Note:
            FAISS不支持直接删除，需要重建索引
        """
        if not FAISS_AVAILABLE or self.index is None:
            return
        
        try:
            original_count = len(self.metadatas)
            
            # 找出需要保留的索引
            indices_to_keep = []
            for i, metadata in enumerate(self.metadatas):
                if metadata.get('document_id') != document_id:
                    indices_to_keep.append(i)
            
            if len(indices_to_keep) == original_count:
                logger.info(f"文档 {document_id} 未找到向量数据")
                return
            
            removed_count = original_count - len(indices_to_keep)
            logger.info(f"开始重建索引，删除文档 {document_id} 的 {removed_count} 个向量...")
            
            # 提取需要保留的向量
            kept_vectors = []
            for i in indices_to_keep:
                try:
                    vec = faiss.vector_to_array(self.index.reconstruct(int(i)))
                    kept_vectors.append(vec)
                except Exception as e:
                    logger.debug(f"跳过索引 {i}: {str(e)}")
            
            # 创建新索引
            self.index = faiss.IndexFlatL2(self.dimension)
            if kept_vectors:
                vectors_array = np.array(kept_vectors, dtype=np.float32)
                self.index.add(vectors_array)
            
            # 更新texts和metadatas
            self.texts = [self.texts[i] for i in indices_to_keep]
            self.metadatas = [self.metadatas[i] for i in indices_to_keep]
            
            logger.info(f"成功删除文档 {document_id} 的 {removed_count} 个向量，剩余 {len(self.texts)} 个")
            
        except Exception as e:
            logger.error(f"删除向量失败: {str(e)}", exc_info=True)
            raise
    
    def merge(self, other_store: 'VectorStore'):
        """
        合并另一个向量库
        
        Args:
            other_store: 另一个向量存储实例
        """
        if not FAISS_AVAILABLE:
            return
        
        try:
            if other_store.count() > 0:
                # 获取其他库的向量
                other_vectors = []
                for i in range(other_store.index.ntotal):
                    vec = other_store.index.reconstruct(int(i))
                    other_vectors.append(vec)
                
                # 添加到当前库
                if other_vectors:
                    vectors = np.array(other_vectors, dtype=np.float32)
                    self.index.add(vectors)
                    self.texts.extend(other_store.texts)
                    self.metadatas.extend(other_store.metadatas)
                    
                    logger.info(f"成功合并 {len(other_vectors)} 个向量")
        except Exception as e:
            logger.error(f"合并向量库失败: {str(e)}")


def load_documents_vector_stores(document_ids: List[str]) -> Optional[VectorStore]:
    """
    加载指定文档的向量存储并合并（新架构：每个文档独立存储）
    
    Args:
        document_ids: 文档ID列表
        
    Returns:
        合并后的向量存储
    """
    if not document_ids or not FAISS_AVAILABLE:
        return None
    
    try:
        from app.core.config import settings
        
        # 自动检测维度
        dimension = get_embedding_dimension(settings.OLLAMA_EMBED_MODEL)
        merged_store = VectorStore(dimension=dimension)
        loaded_count = 0
        
        for doc_id in document_ids:
            # 构建向量库路径（按文档ID存储）
            vector_path = os.path.join(settings.VECTOR_STORE_DIR, f"{doc_id}.faiss")
            
            # 检查 .index 文件是否存在（save 方法会生成 .faiss.index 和 .faiss.meta）
            if os.path.exists(f"{vector_path}.index"):
                try:
                    temp_store = VectorStore(dimension=dimension)
                    temp_store.load(vector_path)
                    
                    # 合并到主库
                    if temp_store.count() > 0:
                        if merged_store.count() == 0:
                            # 第一个向量库，直接使用
                            merged_store = temp_store
                        else:
                            # 后续向量库，需要合并
                            vectors = []
                            for i in range(temp_store.index.ntotal):
                                vec = temp_store.index.reconstruct(int(i))
                                vectors.append(vec)
                            
                            if vectors:
                                vectors_array = np.array(vectors, dtype=np.float32)
                                merged_store.index.add(vectors_array)
                                merged_store.texts.extend(temp_store.texts)
                                merged_store.metadatas.extend(temp_store.metadatas)
                        
                        loaded_count += 1
                        logger.info(f"已加载文档 {doc_id[:8]}... 的向量，共 {temp_store.count()} 条")
                except Exception as e:
                    logger.warning(f"加载文档 {doc_id[:8]}... 的向量失败: {str(e)}")
            else:
                logger.warning(f"文档 {doc_id[:8]}... 的向量文件不存在: {vector_path}")
        
        if merged_store.count() > 0:
            logger.info(f"成功合并 {loaded_count}/{len(document_ids)} 个文档的向量，共 {merged_store.count()} 个向量")
            print(f"✅ 成功加载 {loaded_count}/{len(document_ids)} 个文档的向量，共 {merged_store.count()} 个向量")
            return merged_store
        else:
            logger.warning(f"所有文档都没有向量数据")
            print(f"⚠️  所有文档都没有向量数据")
            return None
    except Exception as e:
        logger.error(f"加载多个文档向量失败: {str(e)}")
        print(f"❌ 加载文档向量失败: {str(e)}")
        return None


def load_multiple_vector_stores(kb_ids: List[str]) -> Optional[VectorStore]:
    """
    【已废弃】加载多个知识库的向量存储并合并（旧架构：按知识库存储）
    
    新架构请使用 load_documents_vector_stores()
    
    Args:
        kb_ids: 知识库ID列表
        
    Returns:
        合并后的向量存储
    """
    logger.warning("load_multiple_vector_stores 已废弃，请使用 load_documents_vector_stores")
    
    if not kb_ids or not FAISS_AVAILABLE:
        return None
    
    try:
        from app.core.config import settings
        
        # 自动检测维度
        dimension = get_embedding_dimension(settings.OLLAMA_EMBED_MODEL)
        merged_store = VectorStore(dimension=dimension)
        
        for kb_id in kb_ids:
            # 构建向量库路径（旧架构）
            vector_path = os.path.join(settings.UPLOAD_DIR, "vectors", kb_id)
            
            if os.path.exists(f"{vector_path}.index"):
                temp_store = VectorStore(dimension=dimension)
                temp_store.load(vector_path)
                
                # 合并到主库
                if temp_store.count() > 0:
                    # 直接添加文本和向量
                    if merged_store.count() == 0:
                        merged_store = temp_store
                    else:
                        # 获取向量并添加
                        vectors = []
                        for i in range(temp_store.index.ntotal):
                            vec = temp_store.index.reconstruct(int(i))
                            vectors.append(vec)
                        
                        if vectors:
                            vectors_array = np.array(vectors, dtype=np.float32)
                            merged_store.index.add(vectors_array)
                            merged_store.texts.extend(temp_store.texts)
                            merged_store.metadatas.extend(temp_store.metadatas)
                    
                    logger.info(f"已加载知识库 {kb_id} 的向量，共 {temp_store.count()} 条")
        
        if merged_store.count() > 0:
            logger.info(f"成功合并 {len(kb_ids)} 个知识库，共 {merged_store.count()} 个向量")
            return merged_store
        else:
            logger.warning(f"所有知识库都没有向量数据")
            return None
    except Exception as e:
        logger.error(f"加载多个向量库失败: {str(e)}")
        return None

